//
//  CFProxySocket.m
//  SdkDemo
//
//  Created by HalloWorld on 2016/12/6.
//  Copyright © 2016年 legendsec. All rights reserved.
//

#import "CFProxySocket.h"
#import <sys/socket.h>
#import <netinet/in.h>
#import <arpa/inet.h>
#import <unistd.h>

void SocketConnectCallBack( CFSocketRef s, CFSocketCallBackType callbackType, CFDataRef address,
                            const void *data, void *info );


typedef void(^pendingSendCompletionBlock)(NSError* error);

@interface CFProxySocket()
{
    CFSocketRef _socketRef;
}
@property (nonatomic, strong, readonly) NSString* destinationHost;
@property (nonatomic, readonly) uint16_t destinationPort;
@property (nonatomic, copy) pendingSendCompletionBlock pendingSendCompletion;
@end

@implementation CFProxySocket

- (instancetype)init
{
    self = [super init];
    if(self)
    {
    }
    return self;
}

- (void)dealloc
{
    _destinationHost = nil;
    self.pendingSendCompletion = nil;
}

- (void)setProxyHost:(NSString *)host port:(uint16_t)port version:(CFProxySocketSOCKSVersion)version {
    _proxyHost = host;
    _proxyPort = port;
    _proxyVersion = version;
}

- (void)setProxyUsername:(NSString *)username password:(NSString *)password {
    _proxyUsername = username;
    _proxyPassword = password;
}

- (BOOL)connectToHost:(NSString*)host onPort:(uint16_t)port
{
    _destinationHost = host;
    _destinationPort = port;
    
    [self releaseSocket];
    
    BOOL isConnect = NO;
    
    CFSocketContext sockContext = { 0, // 结构体的版本，必须为0
        (__bridge void *)(self),  // 一个任意指针的数据，可以用在创建时CFSocket对象相关联。这个指针被传递给所有的上下文中定义的回调。
        NULL,                     // 一个定义在上面指针中的retain的回调， 可以为NULL
        NULL, NULL};

    _socketRef = CFSocketCreate(kCFAllocatorDefault, // 为新对象分配内存，可以为nil
                           PF_INET, // 协议族，如果为0或者负数，则默认为PF_INET
                           SOCK_STREAM, // 套接字类型，如果协议族为PF_INET,则它会默认为SOCK_STREAM
                           IPPROTO_TCP, // 套接字协议，如果协议族是PF_INET且协议是0或者负数，它会默认为IPPROTO_TCP
                           kCFSocketConnectCallBack, // 触发回调函数的socket消息类型，具体见Callback Types
                           SocketConnectCallBack, // 上面情况下触发的回调函数
                           &sockContext // 一个持有CFSocket结构信息的对象，可以为nil
                           );
    
    if (_socketRef != nil)
    {
        // IPV4
        NSString* connectHost = nil;
        NSInteger connectPort = 0;
        
        if(_proxyHost && _proxyPort > 0)
        {
            connectHost = _proxyHost;
            connectPort = _proxyPort;
        }
        else
        {
            connectHost = _destinationHost;
            connectPort = _destinationPort;
        }
        
        struct sockaddr_in addr4;
        memset(&addr4, 0, sizeof(addr4));
        addr4.sin_len = sizeof(addr4);
        addr4.sin_family = AF_INET;
        addr4.sin_port = htons(connectPort);
        addr4.sin_addr.s_addr = inet_addr([connectHost UTF8String]);
        CFDataRef address = CFDataCreate(kCFAllocatorDefault, (UInt8 *)&addr4, sizeof(addr4));
        
//        struct sockaddr_in6 sa6;
//        memset(&sa6, 0, sizeof(sa6));
//        sa6.sin6_port = htons(proxyPort);
//        inet_pton(AF_INET6, [connectHost UTF8String], &(sa6.sin6_addr));
//        CFDataRef address = CFDataCreate(kCFAllocatorDefault, (UInt8 *)&sa6, sizeof(sa6));
        
        CFSocketConnectToAddress(_socketRef,
                                 address,
                                 -1  // 连接超时时间，如果为负，则不尝试连接，而是把连接放在后台进行，如果_socket消息类型为kCFSocketConnectCallBack，将会在连接成功或失败的时候在后台触发回调函数
                                 );
        
        CFRunLoopRef cRunRef = CFRunLoopGetCurrent();
        CFRunLoopSourceRef sourceRef = CFSocketCreateRunLoopSource(kCFAllocatorDefault, _socketRef, 0);
        CFRunLoopAddSource(cRunRef,
                           sourceRef,
                           kCFRunLoopCommonModes
                           );
        CFRelease(sourceRef);

        isConnect = YES;
    }
    
    return isConnect;
}

- (void)releaseSocket
{
    if(_socketRef)
    {
        CFRelease(_socketRef);
        _socketRef = nil;
    }
}

- (void)readStreamData
{
    /**
     int recv( SOCKET s, char FAR *buf, int len, int flags );
     
     不论是客户还是服务器应用程序都用recv函数从TCP连接的另一端接收数据。
     （1）第一个参数指定接收端套接字描述符；
     （2）第二个参数指明一个缓冲区，该缓冲区用来存放recv函数接收到的数据；
     （3）第三个参数指明buf的长度；
     （4）第四个参数一般置0。
     */
    
    int8_t buffer[1024];
    memset(buffer, 0, sizeof(buffer));
    
    long readDataLen;
    //若无错误发生，recv()返回读入的字节数。如果连接已中止，返回0。如果发生错误，返回-1
    while((readDataLen = recv(CFSocketGetNative(_socketRef), buffer, sizeof(buffer), 0)))
    {
        if(readDataLen > 0)
        {
            NSData* readData = [NSData dataWithBytes:buffer length:readDataLen];
            memset(buffer, 0, sizeof(buffer));
            
            if(self.socketReadData)
            {
                self.socketReadData(readData);
            }
        }
    }
    
    if(readDataLen == 0)
    {
        dispatch_async(dispatch_get_main_queue(), ^{
            if(self.socketConnectError)
            {
                int tmpErrorCode = errno;
                NSString* errMsg = [NSString stringWithUTF8String:strerror(tmpErrorCode)];
                NSError* error = [NSError errorWithDomain:errMsg
                                                     code:tmpErrorCode userInfo:nil];
                self.socketConnectError(error);
            }
            [self releaseSocket];
        });
    }
}

- (void)sendStreamData:(NSData*)data
{
    const void* bytes = [data bytes];
    NSUInteger bytesLen = [data length];
    
    //成功则返回实际传送出去的字符数, 失败返回-1
    ssize_t sendData = send(CFSocketGetNative(_socketRef), bytes, bytesLen + 1, 0);
    
    if(sendData < 0)
    {
        dispatch_async(dispatch_get_main_queue(), ^{
            if(self.pendingSendCompletion)
            {
                int tmpErrorCode = errno;
                NSString* errMsg = [NSString stringWithUTF8String:strerror(tmpErrorCode)];
                NSError* error = [NSError errorWithDomain:errMsg
                                                     code:tmpErrorCode userInfo:nil];
                self.pendingSendCompletion(error);
            }
        });
    }
    else
    {
        dispatch_async(dispatch_get_main_queue(), ^{
            if(self.pendingSendCompletion)
            {
                self.pendingSendCompletion(nil);
            }
        });
    }
}

- (void)sendStreamData:(NSData*)data completion:(void(^)(NSError* error))completion
{
    if(self.pendingSendCompletion)
    {
        self.pendingSendCompletion = nil;
    }
    self.pendingSendCompletion = completion;
    
    [self performSelectorInBackground:@selector(sendStreamData:) withObject:data];
}

- (void)proxyConnectErrorMsg:(NSString*)errMsg errorCode:(NSInteger)errorCode
{
    dispatch_async(dispatch_get_main_queue(), ^{
        if(self.socketConnectError)
        {
            NSError* error = [NSError errorWithDomain:errMsg
                                                 code:errorCode userInfo:nil];
            self.socketConnectError(error);
        }
        [self releaseSocket];
    });
}

-(void)proxySOCKS5Connect
{
    char buf[256] = {0};
    buf[0] = '\x05';
    buf[1] = '\x01';
    buf[2] = '\x00';
    
    if ( -1 == send(CFSocketGetNative(_socketRef), buf, 3, 0) )
    {
        int tmpErrorCode = errno;
        [self proxyConnectErrorMsg:@"SOCKS5 Send Hello Error" errorCode:tmpErrorCode];
        return;
    }
    
    if ( 2 != recv(CFSocketGetNative(_socketRef), buf, 2, 0) || '\x05'!=buf[0] )
    {
        int tmpErrorCode = errno;
        [self proxyConnectErrorMsg:@"SOCKS5 Recv Hello Error" errorCode:tmpErrorCode];
        return;
    }
    
    buf[0] = '\x01';
    buf[1] = '\x01';
    buf[2] = '\x00';
    
    size_t pos = 0;
    struct in6_addr dst6;
    if ( 1==inet_pton(AF_INET, [_destinationHost UTF8String], &dst6) )
    {
        // IPv4
        buf[3] = '\x01';
        struct in_addr* dst4 = (struct in_addr*)&dst6;
        memcpy(&buf[4], &dst4->s_addr, 4);
        pos = 8;
    }
    else if ( 1==inet_pton(AF_INET6, [_destinationHost UTF8String], &dst6) )
    {
        // IPv6 not supports
        buf[3] = '\x04';
        memcpy(&buf[4], &dst6.s6_addr, 16);
        pos = 20;
    }
    else
    {
        // hostname
        const char * dstHost = [_destinationHost UTF8String];

        buf[3] = '\x03';
        unsigned char hlen = (unsigned char)strlen(dstHost);
        buf[4] = hlen;
        memcpy(&buf[5], dstHost, hlen);
        pos = 5 + hlen;
        
    }
    uint16_t nport = htons(_destinationPort);
    memcpy(&buf[pos], &nport, 2);
    if ( -1 == send(CFSocketGetNative(_socketRef), buf, pos+2, 0) )
    {
        int tmpErrorCode = errno;
        [self proxyConnectErrorMsg:@"SOCKS5 Send Handshake Error" errorCode:tmpErrorCode];
        return;
    }
    
    memset(buf, 0, sizeof(buf));

    if ( 10 != recv(CFSocketGetNative(_socketRef), buf, 10, 0) || '\x05'!=buf[0] )
    {
        int tmpErrorCode = errno;
        [self proxyConnectErrorMsg:@"SOCKS5 Recv Handshake Error" errorCode:tmpErrorCode];
        return;
    }
    
    if(self.socketDidConnected)
    {
        self.socketDidConnected(self.destinationHost, self.destinationPort);
    }
    [self performSelectorInBackground:@selector(readStreamData) withObject:nil];
}

@end

void SocketConnectCallBack( CFSocketRef s, CFSocketCallBackType callbackType, CFDataRef address,
                            const void *data, void *info )
{
    CFProxySocket* client = (__bridge CFProxySocket *)info;
    
    if (data != NULL)
    {
        NSLog(@"连接失败");
        [client performSelector:@selector(releaseSocket) withObject:nil];
    }
    else
    {
        NSLog(@"连接成功");
        if(client.proxyHost && client.proxyPort > 0)
        {
            [client performSelectorInBackground:@selector(proxySOCKS5Connect) withObject:nil];
        }
        else
        {
            if(client.socketDidConnected)
            {
                client.socketDidConnected(client.destinationHost, client.destinationPort);
            }
            [client performSelectorInBackground:@selector(readStreamData) withObject:nil];
        }
    }
}
